#include "plugin.h"
#include "graph.h"
#include "graph.priv.h"
#include <stack>
#include <list>
#include <vector>
#include <map>
#include <algorithm>
#include <unordered_set>
#include <queue>
#include <array>
#include <stdexcept>
#include <limits>
#ifndef WIN32
#include <unistd.h>
#endif // !WIN32
struct SwcNode {
	size_t next_collision{SIZE_MAX};
	size_t adj_idx{SIZE_MAX};
	unsigned int nadj{0};
	Point point{};
};

struct OctreeNode {
	size_t idx{SIZE_MAX}; // for SwcNode or OctreeNode
	std::array<int32_t, 6> bounds{{0,}}; // left open, right close: (]
	unsigned int cnt{0}; // 0, no branches; >0 no. branches
};

#ifdef BUILD_MINGW
#define INPUT_FILE "C:\\retrace-input.fnt"
#else
#define INPUT_FILE "/tmp/retrace-input.fnt"
#endif

struct AlgFindRetrace final: public AlgFind {
	static AlgFind* creater() {
		return new AlgFindRetrace{};
	}
	// params
	void setParams(const std::vector<ParamValue>& params) override {
		length=params[0].getFloat();
	}
	std::vector<ParamValue> getParams() const override {
		std::vector<ParamValue> vl;
		vl.push_back(ParamValue{length});
		return vl;
	}
	AlgFindRetrace():
		AlgFind(), length{80.0}, octree{}, nodes{}, adjs{}, prev_idx{SIZE_MAX}
	{
	}

	void nodes_add(const SwcNode& node, std::array<int32_t, 6>& bounds) {
		nodes.push_back(node);
		auto& pt=node.point;
		if(pt._x<=bounds[0])
			bounds[0]=pt._x-1;
		if(pt._x>bounds[1])
			bounds[1]=pt._x;
		if(pt._y<=bounds[2])
			bounds[2]=pt._y-1;
		if(pt._y>bounds[3])
			bounds[3]=pt._y;
		if(pt._z<=bounds[4])
			bounds[4]=pt._z-1;
		if(pt._z>bounds[5])
			bounds[5]=pt._z;
	}

	bool inside(const SwcNode& node, const OctreeNode& ot_node) {
		auto& pt=node.point;
		if(pt._x<=ot_node.bounds[0])
			return false;
		if(pt._x>ot_node.bounds[1])
			return false;
		if(pt._y<=ot_node.bounds[2])
			return false;
		if(pt._y>ot_node.bounds[3])
			return false;
		if(pt._z<=ot_node.bounds[4])
			return false;
		if(pt._z>ot_node.bounds[5])
			return false;
		return true;
	}

	void octree_add(size_t idx, SwcNode* node) {
		auto ot_node=&octree[0];
		do {
			if(ot_node->cnt==0) {
				if(ot_node->idx==SIZE_MAX) {
					ot_node->idx=idx;
					ot_node=nullptr;
				} else {
					auto nx=ot_node->bounds[1]-ot_node->bounds[0];
					auto ny=ot_node->bounds[3]-ot_node->bounds[2];
					auto nz=ot_node->bounds[5]-ot_node->bounds[4];
					if(nx<=1 && ny<=1 && nz<=1) {
						node->next_collision=ot_node->idx;
						printMessage("collision...");
						ot_node->idx=idx;
						ot_node=nullptr;
					} else {
						std::array<OctreeNode, 8> branches{};
						size_t j=0;
						for(int i=0; i<8; i++) {
							branches[j].bounds[0]=(i&4)?(ot_node->bounds[0]+nx/2):ot_node->bounds[0];
							branches[j].bounds[1]=(i&4)?ot_node->bounds[1]:(ot_node->bounds[0]+nx/2);
							branches[j].bounds[2]=(i&2)?(ot_node->bounds[2]+ny/2):ot_node->bounds[2];
							branches[j].bounds[3]=(i&2)?ot_node->bounds[3]:(ot_node->bounds[2]+ny/2);
							branches[j].bounds[4]=(i&1)?(ot_node->bounds[4]+nz/2):ot_node->bounds[4];
							branches[j].bounds[5]=(i&1)?ot_node->bounds[5]:(ot_node->bounds[4]+nz/2);
							bool empty{false};
							for(int k=0; k<6; k+=2) {
								if(branches[j].bounds[k+1]-branches[j].bounds[k]<=0)
									empty=true;
							}
							if(!empty) {
								j++;
							}
						}

						auto prev_idx=ot_node->idx;
						bool added{false};
						size_t ot_idx_next{SIZE_MAX};
						for(size_t k=0; k<j; k++) {
							if(inside(nodes[prev_idx], branches[k])) {
								if(added)
									throwError("Overlapping nodes. 1");
								branches[k].idx=prev_idx;
								added=true;
							}
							if(inside(*node, branches[k])) {
								if(ot_idx_next!=SIZE_MAX)
									throwError("Overlapping nodes. 2");
								ot_idx_next=k;
							}
						}
						if(!added)
							throwError("Point not covered. 1 ");
						if(ot_idx_next==SIZE_MAX)
							throwError("Point not covered. 2");

						ot_node->idx=octree.size();
						ot_node->cnt=j;
						ot_idx_next+=octree.size();

						octree.insert(octree.end(), branches.begin(), branches.begin()+j);
						ot_node=&octree[ot_idx_next];
					}
				}
			} else {
				auto ot_idx_max=ot_node->idx+ot_node->cnt;
				OctreeNode* ot_node_next{nullptr};
				for(size_t ot_idx=ot_node->idx; ot_idx<ot_idx_max; ot_idx++) {
					if(inside(*node, octree[ot_idx])) {
						if(ot_node_next)
							throwError("Overlapping nodes. 3");
						ot_node_next=&octree[ot_idx];
					}
				}
				if(!ot_node_next)
					throwError("Point not covered. 3");
				ot_node=ot_node_next;
			}
		} while(ot_node);
	}
	void load_input(const char* fn) {
		Graph gr=GraphPriv::alloc();
		auto gp=GraphPriv::get(gr);
		gp->loadFnt(fn);
		gp->updateModel();

		std::array<int32_t, 6> global_bounds;
		for(int i=0; i<6; i+=2)
			global_bounds[i]=(std::numeric_limits<int32_t>::max)();
		for(int i=1; i<6; i+=2)
			global_bounds[i]=(std::numeric_limits<int32_t>::min)();

		// verts
		printMessage("add nodes...");
		for(auto v: gr.vertices()) {
			SwcNode node{};
			node.nadj=v.neighbors().size();
			node.adj_idx=adjs.size();
			adjs.insert(adjs.end(), node.nadj, SIZE_MAX);
			const Point* pt{nullptr};
			for(auto& p: v.neighbors()) {
				if(p.first.parentVertex()!=v) {
					pt=&p.first.points()[p.second?p.first.points().size()-1:0];
				}
			}
			if(!pt) {
				auto& p=v.neighbors()[0];
				pt=&p.first.points()[p.second?p.first.points().size()-1:0];
			}
			node.point=*pt;
			nodes_add(node, global_bounds);
		}

		// edges
		for(auto e: gr.edges()) {
			size_t lft_idx=e.leftVertex().index();
			size_t rgt_idx=e.rightVertex().index();
			size_t start=0;
			size_t np=e.points().size();
			if(np<2)
				throwError("Edge length should be >=2.");
			if(e.points()[np-1].distTo(nodes[rgt_idx].point)<.05)
				np--;
			if(e.points()[0].distTo(nodes[lft_idx].point)<.05) {
				start++;
				np--;
			}
			if(np==0) {
				adjs[nodes[lft_idx].adj_idx++]=rgt_idx;
				adjs[nodes[rgt_idx].adj_idx++]=lft_idx;
			} else {
				adjs[nodes[lft_idx].adj_idx++]=nodes.size();
				adjs[nodes[rgt_idx].adj_idx++]=nodes.size()+np-1;
				for(size_t i=0; i<np; i++) {
					SwcNode node{};
					node.nadj=2;
					node.adj_idx=adjs.size()+2;
					if(i==0) {
						adjs.push_back(lft_idx);
					} else {
						adjs.push_back(nodes.size()-1);
					}
					if(i==np-1) {
						adjs.push_back(rgt_idx);
					} else {
						adjs.push_back(nodes.size()+1);
					}
					node.point=e.points()[start+i];
					nodes_add(node, global_bounds);
				}
			}
		}

		// check
		{
			size_t adj_idx=0;
			for(auto& node: nodes) {
				if(node.adj_idx!=adj_idx+node.nadj)
					throwError("Check failed. ", node.adj_idx, ' ', adj_idx);
				std::swap(node.adj_idx, adj_idx);
			}
			for(auto i: adjs) {
				if(i==SIZE_MAX)
					throwError("Check failed. 2");
			}
		}

		printMessage("build octree... ");
		OctreeNode ot_root{};
		ot_root.bounds=global_bounds;
		octree.push_back(ot_root);
		for(size_t i=0; i<nodes.size(); i++) {
			octree_add(i, &nodes[i]);
		}
		printMessage("done");
		GraphPriv::free(gr);
	}

	// data
	double length;
	std::vector<OctreeNode> octree;
	std::vector<SwcNode> nodes;
	std::vector<size_t> adjs;

	size_t prev_idx;
	Position prev_pos;
	std::vector<Path> found_paths;

	bool intersect(OctreeNode* ot_node, const Point& p, int32_t l) {
		if(ot_node->bounds[0]>=p._x+l || ot_node->bounds[1]<p._x-l)
			return false;
		if(ot_node->bounds[2]>=p._y+l || ot_node->bounds[3]<p._y-l)
			return false;
		if(ot_node->bounds[4]>=p._z+l || ot_node->bounds[5]<p._z-l)
			return false;
		return true;
	}
	void get_swc_idxes(OctreeNode* ot_node, const Point& p, int32_t l, std::vector<size_t>* ret) {
		// doable
		if(ot_node->cnt==0) {
			auto idx=ot_node->idx;
			while(idx!=SIZE_MAX) {
				ret->push_back(idx);
				idx=nodes[idx].next_collision;
			}
		} else {
			for(unsigned int i=0; i<ot_node->cnt; i++) {
				auto ot_node_br=&octree[ot_node->idx+i];
				if(intersect(ot_node_br, p, l))
					get_swc_idxes(ot_node_br, p, l, ret);
			}
		}
	}
	std::vector<size_t> get_swc_idxes(OctreeNode* ot_node, const Point& p, double l) {
		std::vector<size_t> ret;
		get_swc_idxes(ot_node, p, l*256, &ret);
		std::vector<size_t> refined;
		for(auto i: ret) {
			if(nodes[i].point.distTo(p)<l)
				refined.push_back(i);
		}
		return refined;
	}

	/*
	void get_seed_positions(Edge edg, int idx, const Point& pt, Graph gr, std::unordered_set<size_t>* seeds) {
		if(edg) {
		} else {
		}
	}
	*/

	std::vector<std::deque<size_t>> collect_paths(const std::vector<size_t>& idxes) {
		std::vector<std::deque<size_t>> paths{};
		std::unordered_set<size_t> idx2flag{};
		for(auto idx: idxes)
			idx2flag.insert(idx);
		for(auto start: idxes) {
			if(nodes[start].nadj>2)
				continue;
			auto iter=idx2flag.find(start);
			if(iter==idx2flag.end())
				continue;
			std::deque<size_t> pth{};
			pth.push_back(start);
			idx2flag.erase(iter);

			for(unsigned int k=0; k<nodes[start].nadj; k++) {
				auto cur=start;
				auto next=adjs[nodes[start].adj_idx+k];
				do {
					auto it_next=idx2flag.find(next);
					if(it_next==idx2flag.end())
						break;
					if(k)
						pth.push_back(next);
					else
						pth.push_front(next);
					auto next_nadj=nodes[next].nadj;
					if(next_nadj>2)
						break;
					idx2flag.erase(it_next);
					if(next_nadj<2)
						break;
					auto pnext=&adjs[nodes[next].adj_idx];
					if(*pnext==cur)
						pnext++;
					cur=next;
					next=*pnext;
				} while(true);
			}
			if(pth.size()>=2)
				paths.push_back(pth);
		}
		return paths;
	}

	std::vector<std::pair<Edge, size_t>> get_seeds(const Position& p, Graph gr) {
		if(!p.edge)
			return {};
		std::vector<std::pair<Edge, size_t>> seeds{};
		std::unordered_set<int> edgidx2flag;
		std::queue<std::pair<Edge, size_t>> unfinished{};
		unfinished.push(std::make_pair(p.edge, p.index));
		while(!unfinished.empty()) {
			auto edge=unfinished.front().first;
			auto start=unfinished.front().second;
			unfinished.pop();
			if(edgidx2flag.find(edge.index())!=edgidx2flag.end())
				continue;
			edgidx2flag.insert(edge.index());
			size_t a=start, b=start;
			while(b<edge.points().size()) {
				if(edge.points()[b].distTo(p.point)>=length)
					break;
				seeds.push_back(std::make_pair(edge, b));
				b++;
			}
			while(a-->0) {
				if(edge.points()[a].distTo(p.point)>=length)
					break;
				seeds.push_back(std::make_pair(edge, a));
			}
			if(b>=edge.points().size()) {
				for(auto ep: edge.rightVertex().neighbors()) {
					if(ep.second)
						unfinished.push(std::make_pair(ep.first, ep.first.points().size()-1));
					else
						unfinished.push(std::make_pair(ep.first, 0));
				}
			}
			if(a>=edge.points().size()) {
				for(auto ep: edge.leftVertex().neighbors()) {
					if(ep.second)
						unfinished.push(std::make_pair(ep.first, ep.first.points().size()-1));
					else
						unfinished.push(std::make_pair(ep.first, 0));
				}
			}
		}
		return seeds;
	}

	void find_paths(const Position& p, Graph gr) {
		found_paths.clear();
		printMessage("find ...", p.point._x, ' ', EdgePriv::get(p.edge));
		auto paths=collect_paths(get_swc_idxes(&octree[0], p.point, length));
		auto seeds=get_seeds(p, gr);

		if(seeds.empty()) {
			std::vector<double> mindists(paths.size(), INFINITY);
			std::vector<size_t> idxes(paths.size());
			for(size_t k=0; k<paths.size(); k++) {
				idxes[k]=k;
				auto& path=paths[k];
				Path pth;
				pth.edge0=Edge{};
				pth.index0=-1;
				pth.edge1=Edge{};
				pth.index1=-1;
				for(auto idx: path) {
					auto& pt=nodes[idx].point;
					auto d=pt.distTo(p.point);
					if(d<mindists[k])
						mindists[k]=d;
					pth.points.push_back(nodes[idx].point);
				}
				found_paths.push_back(pth);
			}
			std::sort(idxes.begin(), idxes.end(), [&mindists](size_t a, size_t b) -> bool { return mindists[a]<mindists[b]; });
			std::vector<Path> rawpaths;
			std::swap(rawpaths, found_paths);
			for(auto i: idxes) {
				found_paths.push_back(rawpaths[i]);
			}
		} else {
			for(auto path: paths) {
				Path pth;
				pth.edge0=Edge{};
				pth.index0=-1;
				pth.edge1=Edge{};
				pth.index1=-1;

				for(auto idx: path) {
					auto& pt=nodes[idx].point;
					double mind=INFINITY;
					auto it_seed=seeds.end();
					for(auto it=seeds.begin(); it!=seeds.end(); it++) {
						auto d=it->first.points()[it->second].distTo(pt);
						if(d<mind) {
							mind=d;
							it_seed=it;
						}
					}
					if(mind<0.75) {
						pth.edge1=it_seed->first;
						pth.index1=it_seed->second;
						pth.points.push_back(pt);

						Path tmppth;
						tmppth.edge0=it_seed->first;
						tmppth.index0=it_seed->second;
						tmppth.edge1=Edge{};
						tmppth.index1=-1;
						tmppth.points.push_back(pt);
						std::swap(tmppth, pth);
						found_paths.push_back(tmppth);
					} else {
						pth.points.push_back(pt);
					}
				}
				found_paths.push_back(pth);
			}
			std::vector<Path> rawpaths;
			std::swap(rawpaths, found_paths);
			bool has_conn=false;
			for(auto& path: rawpaths) {
				if(path.points.size()<2)
					continue;
				if(path.points.size()==2) {
					if(path.edge1 && path.edge0)
						continue;
				}

				bool reorder=false;
				if(path.edge0) {
					if(path.edge1) {
						if(path.edge0.points()[path.index0].distTo(p.point)
								> path.edge1.points()[path.index1].distTo(p.point))
							reorder=true;
					}
					has_conn=true;
				} else {
					if(path.edge1) {
						reorder=true;
						has_conn=true;
					} else {
						if(path.points.front().distTo(p.point)
								> path.points.back().distTo(p.point))
							reorder=true;
					}
				}
				if(reorder) {
					Path tmppth;
					tmppth.edge0=path.edge1;
					tmppth.index0=path.index1;
					tmppth.edge1=path.edge0;
					tmppth.index1=path.index0;
					for(size_t i=path.points.size(); i-->0;) {
						tmppth.points.push_back(path.points[i]);
					}
					found_paths.push_back(tmppth);
				} else {
					found_paths.push_back(path);
				}
			}
			rawpaths.clear();
			std::swap(rawpaths, found_paths);
			for(auto& path: rawpaths) {
				if(!path.edge0 && !path.edge1 && has_conn)
					continue;
				found_paths.push_back(path);
			}
			std::sort(found_paths.begin(), found_paths.end(), [p](const Path& a, const Path& b) -> bool { auto da0=a.points.front().distTo(p.point);
			 auto da1=a.points.back().distTo(p.point);
			 if(da1<da0) da0=da1;
			 auto db0=b.points.front().distTo(p.point);
			 auto db1=b.points.back().distTo(p.point);
			 if(db1<db0) db0=db1;
			 return da0<db0;});
		}
			/*

				//pth.points.push_back(p.point);
				//pth.points.push_back(nodes[idx].point);

		//////////////
			//if(nodes[idx].point.distTo(p.point)>length/2) {
				Path pth;

				//pth.edge0=p.edge;
				//pth.index0=p.index;
				pth.edge0=Edge{};
				pth.index0=-1;
				pth.edge1=Edge{};
				pth.index1=-1;

				for(auto idx: path)
					pth.points.push_back(nodes[idx].point);
				//pth.points.push_back(p.point);
				//pth.points.push_back(nodes[idx].point);
				found_paths.push_back(pth);
			//}
		}
		*/

		prev_pos=p;
		prev_idx=0;
		/*
		auto ot_node=get_container(p.point, length);
		auto seeds=get_seed_positions(p, gr);
		auto nearby_nodes=get_swc_nodes(ot_node, p.point, length);
		for(auto seed: seeds) {
			auto very_close_nodes=get_swc_nodes(ot_node, seed.point, 0.05);
			for(...) {
				nearby_nodes2seeds[...]=seed;
			}
		}
		for(
		///////////////
		////////////////
		///////////

		if(seeds.

		pth->
		pth->edge0=cores[store[hit].root].e;
		pth->index0=cores[store[hit].root].ei;
		pth->edge1=Edge{};
		pth->index1=-1;
		p.point);
		*/
	}
	bool check_paths(const Position& p, Graph gr) {

		if(prev_idx>=found_paths.size())
		  return false;
		if(prev_pos.edge!=p.edge)
			return false;
		if(p.edge) {
			if(prev_pos.index!=p.index)
				return false;
		}
		if(!(p.point==prev_pos.point))
			return false;

		for(auto& path: found_paths) {
			bool l=false, r=false;
			for(auto e: gr.edges()) {
				if(path.edge0==e && path.index0<e.points().size())
					l=true;
				if(path.edge1==e && path.index1<e.points().size())
					r=true;
			}
			if(path.edge0 && !l)
				return false;
			if(path.edge1 && !r)
				return false;
		}


		printMessage("skip...");
		return true;
	}
	void find(const Position& p, Graph gr, Path* pth) override {
		if(octree.empty())
			load_input(INPUT_FILE);

		if(!check_paths(p, gr))
			find_paths(p, gr);

		if(prev_idx>=found_paths.size()) {
			pth->points.clear();
			return;
		}
		
		*pth=found_paths[prev_idx++];
		printMessage("edge: ", EdgePriv::get(pth->edge0));






		/*
		double x=p.point.x();
		double y=p.point.y();
		double z=p.point.z();
		idxes.clear();
		for(double dz=-20; dz<20; dz+=2) {
			for(double dy=-20; dy<20; dy+=2) {
				for(double dx=-20; dx<20; dx+=2) {
					auto i=xyzf2index<0>(x+dx, y+dy, z+dz);
					if(i>=0) {
						if(xfunc<0>(cube<0>(i))>.5) {
							idxes.push_back(i);
						}
					}
				}
			}
		}
		//
		//
		while(sync(1000)) {
			usleep(1000);
		}
		pth->points.clear();
		*/
	}
	std::vector<int32_t> getUpdate() const override {
		/*
		std::vector<int32_t> data{};
			for(auto i: idxes) {
			double x, y, z;
			index2xyzf<0>(i, &x, &y, &z);
			Point p(x, y, z);
			data.push_back(p._x);
			data.push_back(p._y);
			data.push_back(p._z);
			}
		return data;
		*/
		return {};
	}
};

static void fnt_load(std::vector<AlgorithmInfo*>& algs, std::vector<GraphReaderInfo*>& readers) {
	auto algfind=new AlgInfo<AlgFind>{"Retrace", 40,
		"Re-trace neurons. Input is assumed to be located at '" INPUT_FILE "'.",
		{
			ParamSpec{"Length", "...", 1.0, 270.0, 80.0, "µm"}
		},
		AlgFindRetrace::creater};
	algs.push_back(algfind);
}

static void fnt_unload() {
}
extern "C"
{
	PLUGIN_EXPORT
		FNT_PLUGIN("Retrace", fnt_load, fnt_unload, "Re-trace neurons.");
}
